%% initialization
ROOTDIR = '..';
MEXBINDIR = [ROOTDIR filesep 'bin'];
[~,~,~] = mkdir(MEXBINDIR);
addpath([ROOTDIR filesep 'common'], MEXBINDIR, [ROOTDIR filesep 'IBP'], fullfile(ROOTDIR, '..', 'daSVM'));

N = 5; D = 5;
alphav = 4.5;
sigmax = 4;
sigmav = 1/2;
sigmaz = 1/3;

poisstrunc = 10;
maxiter = 500;
numiter = 500;
algtype = 1;

gfuncs = {
    @(X,Z,V)size(Z,2), ...
    @(X,Z,V)mean(sum(Z~=0,2)), ...
    @(X,Z,V)sum(sum(Z~=0) == 1), ...
    @(X,Z,V)sum(sum(Z~=0) == (N+mod(N,2))/2), ...
    @(X,Z,V)max([2.^(N-1:-1:0)*(Z~=0), 0]), ...
    @(X,Z,V)median([2.^(N-1:-1:0)*(Z~=0), 0]), ...
    @(X,Z,V)numel(unique(2.^(N-1:-1:0)*(Z~=0))), ...
    @(X,Z,V)mean([V(:);0]), ...
    @(X,Z,V)max(std(V(:), 1), 0), ...
    @(X,Z,V)fobj(X, Z, V, alphav, sigmax, sigmaz, sigmav) };

gfuncnames = {
    '# of features', ...
    'avg. # of active features by each object', ...
    '# of singleton features', ...
    '# of half-full features', ...
    'maximum history', ...
    'median history', ...
    '# of unique histories', ...
    'mean(V)', ...
    'std(V,1)', ...
    'log joint density' };

gavgvals = [
    alphav*hrmsum(N), ...
    alphav, ...
    alphav, ...
    alphav*2/(N+mod(N,2)), ...
    nan, ...
    nan, ...
    nan, ...
    0, ...
    sigmav, ...
    nan ];

gstdvals = [
    sqrt(alphav*hrmsum(N)), ...
    nan, ...
    sqrt(alphav), ...
    sqrt(alphav*2/(N+mod(N,2))), ...
    nan, ...
    nan, ...
    nan, ...
    nan, ...
    nan, ...
    nan ];

mrgcndstats = zeros(numiter, numel(gfuncs));
succndstats = zeros(size(mrgcndstats));
avgmrgcndstats = zeros(size(mrgcndstats));
stdmrgcndstats = zeros(size(mrgcndstats));
avgsuccndstats = zeros(size(succndstats));
stdsuccndstats = zeros(size(succndstats));

savedir = input('save to: ', 's');
mkdir(savedir);

%% getting-it-right
Z2 = ribprnd(N, alphav, @(sz)sigmaz*randn(sz));
V2 = sigmav*randn(D, size(Z2, 2));
for i = 1:numiter
    Z1 = ribprnd(N, alphav, @(sz)sigmaz*randn(sz));
    V1 = sigmav*randn(D, size(Z1, 2));
    X1 = Z1*V1' + sigmax*randn(N, D);
    mrgcndstats(i,:) = cellfun(@(f)f(X1,Z1,V1), gfuncs);
    
    X2 = Z2*V2' + sigmax*randn(N, D);
    for j = 1:maxiter
        [Z2,V2] = mcmc(X2, Z2, V2, alphav, sigmax, sigmaz, sigmav, ...
            poisstrunc, algtype, -1);
    end
    if algtype == 3
        afidx = any(Z2);
        Z = Z2(:,afidx);
        V = V2(:,afidx);
    elseif algtype == 1
        Z = Z2; V = V2;
    end
    succndstats(i,:) = cellfun(@(f)f(X2,Z,V), gfuncs);
    
    avgmrgcndstats(i,:) = mean(mrgcndstats(1:i,:),1);
    stdmrgcndstats(i,:) = std(mrgcndstats(1:i,:),1,1);
    avgsuccndstats(i,:) = mean(succndstats(1:i,:),1);
    stdsuccndstats(i,:) = std(succndstats(1:i,:),1,1);
    
    stats = [
        num2str([avgmrgcndstats(i,:)', stdmrgcndstats(i,:)'], ...
                ['%.3f',char(177),'%.3f']), ...
        repmat(blanks(4), numel(gfuncs), 1), ...
        num2str([avgsuccndstats(i,:)', stdsuccndstats(i,:)'], ...
                ['%.3f',char(177),'%.3f'])
    ];
    fprintf('%d:\n', i);
    disp(stats);
end

save([savedir filesep 'stats.mat'], 'N', 'D', 'alphav', 'sigmaz', 'sigmav', 'sigmax', ...
    'poisstrunc', 'maxiter', 'algtype', 'numiter', ...
    'gfuncs', 'gfuncnames', 'gavgvals', 'gstdvals', ...
    'mrgcndstats', 'avgmrgcndstats', 'stdmrgcndstats', ...
    'succndstats', 'avgsuccndstats', 'stdsuccndstats');

%% display the results
hyperparamnames = {'alpha', 'sigmax', 'sigmaz', 'sigmav'};
hyperparamvals = [alphav, sigmax, sigmaz, sigmav];
disp([char(hyperparamnames), ...
    repmat(' = ', numel(hyperparamnames), 1), ...
    num2str(hyperparamvals', '%.4f')]);

for j = 1:numel(gfuncs)
    if isnan(gavgvals(j))
        plot(avgmrgcndstats(:,j)+stdmrgcndstats(:,j), '-.r'); hold on;
        plot(avgmrgcndstats(:,j)-stdmrgcndstats(:,j), '-.r');
        plot(avgmrgcndstats(:,j), '-.b');
        plot(avgsuccndstats(:,j)+stdsuccndstats(:,j),'r');
        plot(avgsuccndstats(:,j)-stdsuccndstats(:,j),'r');
        plot(avgsuccndstats(:,j),'b'); hold off;
        title([gfuncnames{j}, ': marginal/successive -conditional']);
        pause
        continue;
    end
    subplot(2,1,1);
    plot(avgmrgcndstats(:,j)+stdmrgcndstats(:,j), 'r'); hold on;
    plot(avgmrgcndstats(:,j)-stdmrgcndstats(:,j), 'r');
    plot(avgmrgcndstats(:,j), 'b');
    plot(1:size(avgmrgcndstats,1), gavgvals(j)+gstdvals(j), '-.r');
    plot(1:size(avgmrgcndstats,1), gavgvals(j)-gstdvals(j), '-.r');
    plot(1:size(avgmrgcndstats,1), gavgvals(j), '-.b'); hold off;
    title([gfuncnames{j}, ': marginal-conditional']);
    
    subplot(2,1,2);
    plot(avgsuccndstats(:,j)+stdsuccndstats(:,j),'r'); hold on;
    plot(avgsuccndstats(:,j)-stdsuccndstats(:,j),'r');
    plot(avgsuccndstats(:,j),'b');
    plot(1:size(avgsuccndstats,1), gavgvals(j)+gstdvals(j), '-.r');
    plot(1:size(avgsuccndstats,1), gavgvals(j)-gstdvals(j), '-.r');
    plot(1:size(avgsuccndstats,1), gavgvals(j), '-.b'); hold off;
    title([gfuncnames{j}, ': successive-conditional']);
    pause
end
close all